import os
import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

# 创建save_model目录
os.makedirs('save_model', exist_ok=True)

# 生成随机训练数据
np.random.seed(42)  # 确保结果可重现
n_samples = 1000
n_features = 5

# 生成特征数据
X = np.random.randn(n_samples, n_features)
feature_names = [f'feature_{i}' for i in range(n_features)]

# 生成目标变量（简单的线性关系加噪声）
y = np.sum(X, axis=1) + np.random.randn(n_samples) * 0.1

# 创建DataFrame
df = pd.DataFrame(X, columns=feature_names)
df['target'] = y

# 分割数据
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 训练随机森林模型
model = RandomForestRegressor(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# 评估模型
y_pred = model.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
print(f"模型均方误差: {mse:.4f}")

# 保存模型和特征信息
model_data = {
    'model': model,
    'feature': feature_names
}

joblib.dump(model_data, 'save_model/model.joblib')
print("模型已保存到 save_model/model.joblib")

# 生成一个测试样本
test_sample = {f'feature_{i}': np.random.randn() for i in range(n_features)}
print(f"\n测试样本: {test_sample}")

# 测试预测
test_df = pd.DataFrame([test_sample])
test_input = test_df[feature_names]
test_input = np.array(test_input)
prediction = model.predict(test_input)
print(f"预测结果: {prediction[0]:.4f}")
